/*
** Copyright (C) 2010 Calin Radoni
**
** This file is part of CHScanner.
**
** CHScanner is free software: you can redistribute it and/or modify
** it under the terms of the GNU General Public License as published by
** the Free Software Foundation, either version 3 of the License, or
** (at your option) any later version.
**
** CHScanner is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
** GNU General Public License for more details.
**
** You should have received a copy of the GNU General Public License
** along with CHScanner. If not, see <http://www.gnu.org/licenses/>.
*/

#include "StdAfx.h"
#include ".\scanmethodwol.h"

//--------------------------------------------------------------------

ScanMethodWOL::ScanMethodWOL(void) : ScanMethod()
{
	myType = __wmtWakeOnLAN;
}

ScanMethodWOL::~ScanMethodWOL(void)
{
}

//--------------------------------------------------------------------

int ScanMethodWOL::InternalInit(void)
{
	CString s1, s2;
	int val;

	FillRandomBuffer();
	
	if(workPhase!=__wptDiscovery) return MSR_ERR_PHASE;

	useUdp  = false;
	useIcmp = false;

	s1 = scanParameters.Trim();
	if(s1.IsEmpty()){
		s1 = _T("40000");
		s2 = _T("U");
	}
	else{
		s2 = s1.Left(1);
		if(s1.GetLength()>1)
			s1 = s1.Mid(1);
		else
			s1 = _T("40000");
	}

	if(s2.CompareNoCase(_T("D"))==0){
		useUdp  = true;
		useIcmp = false;
		s1 = _T("7");
	}
	if(s2.CompareNoCase(_T("N"))==0){
		useUdp  = true;
		useIcmp = false;
		s1 = _T("9");
	}
	if(s2.CompareNoCase(_T("U"))==0){
		useUdp  = true;
		useIcmp = false;
	}

	if(s2.CompareNoCase(_T("I"))==0){
		useUdp  = false;
		useIcmp = true;
	}

	if(!useUdp && !useIcmp)
		return MSR_ERR_PARAMETER;

	val = _tstoi(s1);
	if(val<0 || val>65535) return MSR_ERR_PARAMETER;
	if(val==0 && useUdp)   return MSR_ERR_PARAMETER;
	wolPort = (WORD)val;

	return MSR_ERR_OK;
}

bool ScanMethodWOL::GetArpCacheRec(ScanPair* spIn, MArpCacheRec* recIn)
{
	if(spIn==NULL) return false;
	if(spIn->addr==NULL) return false;

	if(recIn==NULL) return false;
	if(arpCacheList==NULL) return NULL;

	if(!arpCacheList->GetRecordByIPv4(spIn->addr->addrIP.GetD(), recIn)) return false;

	return true;
}

int ScanMethodWOL::InternalExecute(void)
{
	MArpCacheRec arpCacheRec;

	if(addrList==NULL) return MSR_ERR_PARAMETER;

	ScanPair sp;
	while(GetNextPair(&sp, __sptIpOnly) && !stopStopStop){
		if(sp.addr!=NULL){
			if(!sp.addr->addrIP.IsEmpty()){
				if(sp.addr->alive){
					// why wake-up an already awake machine ?
				}
				else{
					if(GetArpCacheRec(&sp, &arpCacheRec)){
						// I have a MAC and an IP so ...
						if(config->xMinDelay>0){
							if(lastScanTime!=0){
								waitDone = false;
								while(!waitDone){
									CheckResponses();
									SleepEx(1, TRUE);
									scanTime = GetTickCount();
									scanTime = ComputeDeltaTime(lastScanTime, scanTime);
									if(scanTime>=config->xMinDelay)
										waitDone = true;
								}
							}
							lastScanTime = GetTickCount();
						}

						SendWolPacket(&arpCacheRec);
					}

					if((packSent % notifTcp) == 0 && packSent>0)
						if(hWnd!=NULL)
							PostMessage(hWnd, CHS_MSG_PACKETS_SENT, 0, (LPARAM)packSent);
				}
			}
		}
	}

	if(hWnd!=NULL)
		PostMessage(hWnd, CHS_MSG_PACKETS_SENT, (WPARAM)1, (LPARAM)packSent);

	return MSR_ERR_OK;
}

//--------------------------------------------------------------------

void ScanMethodWOL::ComputeRandomValues(void)
{
	// nothing to do
}

//--------------------------------------------------------------------

bool ScanMethodWOL::SendWolPacket(MArpCacheRec* recIn)
{
	BYTE srcMac[6];
	BYTE dstMac[6];
	BYTE wolMac[6];
	DWORD sa, da;
	WORD dp;
	bool res;

	winPcap->FillSrcMAC(&srcMac[0]);	// nu conteaza !!!

	// TODO try broadcast, unicast or multicast address here 
	// = quit a nonsens if you do know the MAC address
	// but, very good if you do NOT know the MAC address
	dstMac[0] = recIn->addrMAC.addr[0];
	dstMac[1] = recIn->addrMAC.addr[1];
	dstMac[2] = recIn->addrMAC.addr[2];
	dstMac[3] = recIn->addrMAC.addr[3];
	dstMac[4] = recIn->addrMAC.addr[4];
	dstMac[5] = recIn->addrMAC.addr[5];	// conteaza
	
	wolMac[0] = recIn->addrMAC.addr[0];
	wolMac[1] = recIn->addrMAC.addr[1];
	wolMac[2] = recIn->addrMAC.addr[2];
	wolMac[3] = recIn->addrMAC.addr[3];
	wolMac[4] = recIn->addrMAC.addr[4];
	wolMac[5] = recIn->addrMAC.addr[5];	// conteaza

	sa = winPcap->GetLocalIP();			// nu conteaza !!!
	da = recIn->addrIP.GetD();			// nu conteaza !!!
	dp = wolPort;						// nu conteaza !!!

	res = false;
	if(useUdp)
		res = packetBuilderEx->BuildPackWolUdp(srcMac, dstMac, sa, da, dp, wolMac);
	if(useIcmp)
		res = packetBuilderEx->BuildPackWolIcmp(srcMac, dstMac, sa, da, wolMac);
	if(!res) return false;

	return SendPacket(packetBuilderEx->GetPacket(), (int)packetBuilderEx->GetPacketSize());
}

//--------------------------------------------------------------------
